import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class Actor(nn.Module):
    def __init__(self, args, agent_id, num_action, env):
        super(Actor, self).__init__()
        self.agent_id = agent_id
        self.args = args
        self.env = env
        
        # 简化输入维度：10维
        # 当前位置(2) + 目标相对位置(2) + 其他UAV相对位置(4) + 历史移动(2) = 10维
        input_dim = 10
        hidden_dim = 32
        
        # 简化网络结构
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_action)
        )
        
        # 简化权重初始化
        self._initialize_weights()
        
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.constant_(m.bias, 0.01)
    
    def get_action_mask(self, uav_pos):
        """生成简单的动作掩码，避免越界"""
        mask = torch.ones(4, device=device)
        x, y = uav_pos
        
        # 边界检查
        if y >= self.args.grid_size - 1:  # 上边界
            mask[0] = 0.1
        if y <= 0:  # 下边界
            mask[1] = 0.1
        if x <= 0:  # 左边界
            mask[2] = 0.1
        if x >= self.args.grid_size - 1:  # 右边界
            mask[3] = 0.1
            
        return mask
    
    def forward(self, x, uav=None, episode=None):
        logits = self.network(x)
        
        # 应用动作掩码
        if uav is not None:
            mask = self.get_action_mask(uav.position)
            masked_logits = logits * mask.unsqueeze(0)
        else:
            masked_logits = logits
            
        return F.softmax(masked_logits, dim=-1)

# 保持原有的Critic网络不变，或者也可以简化
class Critic(nn.Module):
    def __init__(self, args):
        super(Critic, self).__init__()
        
        # 输入：所有智能体的状态 + 动作
        state_dim = 10 * args.num_uavs  # 每个智能体10维状态
        action_dim = 4 * args.num_uavs  # 每个智能体4维动作
        input_dim = state_dim + action_dim
        
        self.network = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )
        
        self._initialize_weights()
        
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.constant_(m.bias, 0.01)
    
    def forward(self, state, action):
        x = torch.cat([state, action], dim=1)
        return self.network(x)